'''
@Company: TWL
@Author: xue jian
@Email: xuejian@kanzhun.com
@Date: 2020-06-29 10:51:50
'''
'''
215. 数组中的第K个最大元素
在未排序的数组中找到第 k 个最大的元素。请注意，你需要找的是数组排序后的第 k 个最大的元素，而不是第 k 个不同的元素。

示例 1:

输入: [3,2,1,5,6,4] 和 k = 2
输出: 5
示例 2:

输入: [3,2,3,1,2,4,5,5,6] 和 k = 4
输出: 4
说明:

你可以假设 k 总是有效的，且 1 ≤ k ≤ 数组的长度。
采用线性查找算法。
'''

from typing import List
class Solution:
    def sort_2_media(self, nums):
        if len(nums)<=5:
            nums.sort(reverse = True)
            # print('neibu', nums)
            nums[0], nums[len(nums)>>1] = nums[len(nums)>>1], nums[0] #寻找中位数并放在数组首位
            # print(nums)
            return 
        else:
            l = (len(nums)+4)//5
            for i in range(l):
                r = min(i*5+5, len(nums))
                tmp = nums[i*5:r]
                self.sort_2_media(tmp)
                nums[i*5:r] = tmp
                # print('waibu', nums[i*5: r])
                nums[i], nums[i*5] = nums[i*5], nums[i] #将中位数并放在数组前排
            # print('haha', nums)
            tmp = nums[:l]
            self.sort_2_media(tmp) 
            nums[:l] = tmp
            

    def partition(self, nums):
        nums[0], nums[-1] = nums[-1], nums[0]
        target = nums[-1]
        i=0
        for j in range(len(nums)-1):
            if nums[j]>= target:
                nums[i], nums[j] = nums[j], nums[i]
                i+=1
        # print(i)
        nums[i], nums[-1] = nums[-1], nums[i]
        return i, nums[i]

    def BFPRT(self, nums, k):
        self.sort_2_media(nums)
        # print(nums)
        i, v = self.partition(nums)
        # print(i, v)
        if i == k:
            return v
        elif i>k:
            # print(i, k)
            return self.BFPRT(nums[:i], k)
        else:
            # print(i, k)
            return self.BFPRT(nums[i:], k-i)
    def findKthLargest(self, nums: List[int], k: int) -> int:
        return self.BFPRT(nums, k-1)

if __name__ == "__main__":
    solution = Solution()
    nums = [3,2,3,1,2,4,5,5,6]
    k = 4
    print(solution.findKthLargest(nums, k))